#include "../BST/BST.hpp"
#include <algorithm>

#ifndef _REDBLACK_HPP_
#define _REDBLACK_HPP_

/* The outside node is also black node. */
#define IsBlack(p) (!(p) || (RB_BLACK == (p)->color))
#define IsRed(p) (!IsBlack((p)))
/* Check if the blakc height of the node x is updated. */
/* If it is updated, the black height of the children of x are equal, */
/* and if x is red, the black height of x is equal to its children, else 
 * the black height of x is equal to the black height of its children plus 1. */
#define BlackHeightUpdated(x) ( \
	(stature((x).lc) == stature((x).rc)) && \
	((x).height == (IsRed(&(x)) ? stature((x).lc) : stature((x).lc) + 1)) \
)

template <class T>
class RedBlack: public BST<T> {
public:
	BinNodePosi(T) insert(const T &e) {
		BinNodePosi(T) &x = search(e);
		if(x) return x;

		/* The default node is red! */
		x = new BinNode<T>(e, _hot, nullptr, nullptr, -1);
		++_size;

		solveDoubleRed(x);
		return x ? x : _hot->parent;
	}

	/* It maybe cause double-black situations. */
	bool remove(const T &e) {
		BinNodePosi(T) &x = search(e);
		if(!x) return false;

		/* r is the node that replaces x. */
		/* The sibling of r is null. */
		BinNodePosi(T) r = removeAt(x, _hot);
		--_size;
		if(!_size) return true;

		/* If the deleted one is ROOT. */
		if(!_hot) {
			_root->color = RB_BLACK;
			updateHeight(_root);
			return true;
		}

		/* x must have parent. */
		/* ------------------------- */

		/* x is red. */
		if(BlackHeightUpdated(*_hot)) {
			return true;
		}

		/* x is black. */
		/* if r is red. */
		if(IsRed(r)) {
			r->color = RB_BLACK;
			r->height++;
			return true;
		}


		/* x and r are both black. */
		/* if r is black, the black height of r->parent will decrease. */
		solveDoubleBlack(r);
		return true;
	}

protected:
	void solveDoubleRed(BinNodePosi(T) x) {
		if(IsRoot(*x)) {
			_root->color = RB_BLACK;
			_root->height++;
			return;
		}

		BinNodePosi(T) p = x->parent;
		if(IsBlack(p)) return; 
		BinNodePosi(T) g = p->parent;
		BinNodePosi(T) u = uncle(x);

		/* RR-1 */
		if(IsBlack(u)) {
			/* p, x are at the same side. */
			if(IsLChild(*x) == IsLChild(*p)) p->color = RB_BLACK;
			/* p, x are at the different side. */
			else x->color = RB_BLACK;
			g->color = RB_RED;
			
			/* connect34 */
			BinNodePosi(T) gg = g->parent;
			BinNodePosi(T) &from_parent_to = FromParentTo(*g);
			BinNodePosi(T) r = from_parent_to = rotateAt(x);
			r->parent = gg;
		}
		/* RR-2 */
		else {
			/* p, u are both red. */
			/* only need to change the color. */
			p->color = RB_BLACK;
			p->height++;
			u->color = RB_BLACK;
			u->height++;

			if(!IsRoot(*g)) g->color = RB_RED;
			solveDoubleRed(g);
		}
	}

	void solveDoubleBlack(BinNodePosi(T) r) {
		BinNodePosi(T) p = r ? r->parent : _hot;
		if(!p) return;
		BinNodePosi(T) s = (r == p->lc ? p->rc : p->lc);

		if(IsBlack(s)) {
			BinNodePosi(T) t = nullptr;
			if(IsRed(s->rc)) t = s->rc;
			if(IsRed(s->lc)) t = s->lc;

			/* BB-1 */
			if(t) {
				RBColor oldColor = p->color;

				BinNodePosi(T) & from_parent_to = FromParentTo(*p);
				BinNodePosi(T) b = from_parent_to = rotateAt(t);
				if(HasLChild(*b)) {
					b->lc->color = RB_BLACK;
					updateHeight(b->lc);
				}
				if(HasRChild(*b)) {
					b->rc->color = RB_BLACK;
					updateHeight(b-rc);
				}

				b->color = oldColor;
				updateHeight(b);
			}
			else {
				s->color = RB_RED;
				s->height--;

				/* BB-2R */
				if(IsRed(p)) {
					p->color = RB_BLACK;
				}
				/* BB-2B */
				else {
					p->height--;
					solveDoubleBlack(p);
				}
			}
		}
		/* BB-3 */
		else {
			s->color = RB_BLACK;
			p->color = RB_RED;

			BinNodePosi(T) t = IsLChild(*s) ? s->lc : s->rc;

			_hot = p;
			BinNodePosi(T) & from_parent_to = FromParentTo(*p);
			from_parent_to = rotateAt(t);

			solveDoubleBlack(r);
		}
	}

	/* compute the black height of x. */
	/* The leaf node is black. */
	/* The black height of all leaves is 0!!! */
	int updateHeight(BinNodePosi(T) x) {
		x->height = std::max(stature(x->lc), stature(x->rc));
		/* If x is black, x should be counted. */
		/* Note: the black height of null is 0, 
		 * though stature(null) == -1, 
		 * it does not affect the comparison between black height. 
		 * */
		if(IsBlack(x)) {
			++x->height;
		}
		return x->height;
	}
};

#endif